package io.gatling.extensions.redis.protocol

import io.gatling.core.config.GatlingConfiguration
import io.gatling.core.protocol.{Protocol, ProtocolComponents, ProtocolKey}
import io.gatling.core.session.Session
import io.gatling.core.{CoreComponents, protocol}
import redis.clients.jedis.Jedis

case class RedisProtocol(host: String, port: Int, password: Option[String]) extends Protocol {
  type Components = RedisComponents
}

object RedisProtocol {

  def apply(host: String, port: Int, password: Option[String] = None): RedisProtocol = new RedisProtocol(host, port, password)

  val redisProtocolKey: ProtocolKey[RedisProtocol, RedisComponents] = new ProtocolKey[RedisProtocol, RedisComponents] {

    override def protocolClass: Class[protocol.Protocol] = classOf[RedisProtocol].asInstanceOf[Class[io.gatling.core.protocol.Protocol]]

    override def defaultProtocolValue(configuration: GatlingConfiguration): RedisProtocol = RedisProtocol("127.0.0.1", 6379)

    override def newComponents(coreComponents: CoreComponents): RedisProtocol => RedisComponents = {
      redisProtocol => RedisComponents(redisProtocol)
    }
  }
}

case class RedisComponents(redisProtocol: RedisProtocol) extends ProtocolComponents {

  val jedisMap: scala.collection.mutable.Map[Long, Jedis] = scala.collection.mutable.HashMap()

  override def onStart: Session => Session = {
    session => {
      println("RedisComponents start: " + session)
      val jedis: Jedis = new Jedis(redisProtocol.host, redisProtocol.port)

      redisProtocol.password match {
        case Some(value) => jedis.auth(value)
        case None =>
      }

      jedis.set(String.valueOf(session.userId), session.toString)

      jedisMap(session.userId) = jedis

      session
    }
  }

  override def onExit: Session => Unit = {
    session => {
      println("RedisComponents exit: " + session)
    }
  }
}

case class RedisProtocolBuilder(host: String, port: Int, password: Option[String]) {
  def build: RedisProtocol = RedisProtocol(host, port, password)
}

object RedisProtocolBuilder {
  implicit def toRedisProtocol(builder: RedisProtocolBuilder): RedisProtocol = builder.build
}